# 'Friendly Flow' is a Server Framework for Python 2.5 and up
# Copyright (C) Xaba Software http://www.xaba.nl/
#
# This file is part of 'Friendly Flow'
#
# 'Friendly Flow' is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# 'Friendly Flow' is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with 'Friendly Flow'; if not, write to the Free Software
# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA

from sys import exc_info
from traceback import format_tb
from types import GeneratorType

class RETURN:
    def __init__(self, value=None):
        self.value = value
    
    def __eq__(self, other):
        return isinstance(other, RETURN) and self.value == other.value
    
    def __repr__(self):
        return "RETURN(%s)" % repr(self.value)

def flatten(generator):
    """Flatten takes a generator which may
        * other generators (will be called)
        * exactly one RETURN statement
        * values (all others)
    and returns a flattened generator with
        * exactly one RETURN statement
        * values (all others)
    """
    lastReturnValue = None
    
    stack = []
    top = generator
    current = generator.next()
    
    while True:
        if isinstance(current, GeneratorType):
            stack.append(generator)
            generator = current
        elif isinstance(current, RETURN):
            lastReturnValue = current.value
            generator.close()
            if not stack:
                yield current
                raise StopIteration
            generator = stack.pop()
        else:
            yield current
        
        try:
            current = generator.send(lastReturnValue)
        except:
            type, e, tb = exc_info()
            while e:
                if not stack:
                    raise type, e, tb
                generator = stack.pop()
                try:
                    if type == StopIteration:
                        current = generator.send(None)
                    else:
                        current = generator.throw(type, e, tb)
                    e = None
                except:
                    type, e, tb = exc_info()
        finally:
            lastReturnValue = None
        
    